/*
 * @lc app=leetcode.cn id=498 lang=golang
 * @lcpr version=30118
 *
 * [498] 对角线遍历
 */

// @lcpr-template-start

// @lcpr-template-end
// @lc code=start
func findDiagonalOrder(mat [][]int) []int {
	n, m := len(mat), len(mat[0])
	if n == 1 {
		return mat[0]
	}
	rightUpper, right2left, leftDown, left2right := true, false, false, false
	res := make([]int, 0, n*m)
	res = append(res, mat[0][0])
	for i, j := 0, 0; i < n && j < m && (rightUpper || right2left || leftDown || left2right); {
		// 向右上遍历
		for i-1 >= 0 && j+1 < m {
			i--
			j++
			rightUpper = true
			res = append(res, mat[i][j])
		}

		// 从右上向左下拐弯
		// 先向右遍历
		if j+1 < m {
			j++
			right2left = true
			res = append(res, mat[i][j])
		} else if i+1 < n {
			// 向右遍历失败向下遍历
			i++
			right2left = true
			res = append(res, mat[i][j])
		} else {
			return res
		}

		// 向左下遍历
		for i+1 < n && j-1 >= 0 {
			i++
			j--
			leftDown = true
			res = append(res, mat[i][j])
		}

		// 从左下向右上拐弯
		// 先向下遍历
		if i+1 < n {
			i++
			left2right = true
			res = append(res, mat[i][j])
		} else if j+1 < m {
			// 向下遍历失败向右遍历
			j++
			left2right = true
			res = append(res, mat[i][j])
		} else {
			return res
		}

	}

	return res
}

// @lc code=end

/*
// @lcpr case=start
// [[1,2,3],[4,5,6],[7,8,9]]\n
// @lcpr case=end

// @lcpr case=start
// [[1,2],[3,4]]\n
// @lcpr case=end

*/

